{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deep Learning Barrier Option\n",
    "\n",
    "We used Numba and CuPy in the previous notebook to run Monte Carlo simulation to determine the price of the Asian Barrier option. A Monte Carlo simulation needs millions of paths to get an accurate answer which is computationally intensive. [Ryan et al (2018)](https://arxiv.org/abs/1809.02233) showed that a deep learning model can be trained to value derivatives. The deep learning model is accurate and very fast, capable of producing valuations a million times faster than traditional models. In the this notebook, we will use a fully connected network to learn the pricing mode of the Asian Barrier option. Monte Carlo simulation is used as pricing ground truth for the training. We use the same Asian Barrier Option model as last notebook with parameters listed as following:\n",
    "\n",
    "```\n",
    "T - Maturity (yrs.)\n",
    "S - Spot (usd)\n",
    "K - Strike (usd)\n",
    "sigma - Volatility (per.)\n",
    "r - Risk Free Rate (per.)\n",
    "mu - Stock Drift Rate (per.)\n",
    "B - Barrier (usd)\n",
    "```\n",
    "\n",
    "### Batched Data generation\n",
    "\n",
    "The dataset is an important part of the Deep learning training. We will modify the previous single Asian Barrier Option pricing code to handle a batch of Barrier Option pricing. \n",
    "\n",
    "Loading all the necessary libraries:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy\n",
    "import numpy as np\n",
    "import math\n",
    "import time\n",
    "import torch\n",
    "cupy.cuda.set_allocator(None)\n",
    "from torch.utils.dlpack import from_dlpack"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The CuPy version of batched barrier option pricing simulation is as follows:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "cupy_batched_barrier_option = cupy.RawKernel(r'''\n",
    "extern \"C\" __global__ void batched_barrier_option(\n",
    "    float *d_s,\n",
    "    const float T,\n",
    "    const float * K,\n",
    "    const float * B,\n",
    "    const float * S0,\n",
    "    const float * sigma,\n",
    "    const float * mu,\n",
    "    const float * r,\n",
    "    const float * d_normals,\n",
    "    const long N_STEPS,\n",
    "    const long N_PATHS,\n",
    "    const long N_BATCH)\n",
    "{\n",
    "  unsigned idx =  threadIdx.x + blockIdx.x * blockDim.x;\n",
    "  unsigned stride = blockDim.x * gridDim.x;\n",
    "  unsigned tid = threadIdx.x;\n",
    "  const float tmp3 = sqrt(T/N_STEPS);\n",
    "\n",
    "\n",
    "  for (unsigned i = idx; i<N_PATHS * N_BATCH; i+=stride)\n",
    "  {\n",
    "    int batch_id = i / N_PATHS;\n",
    "    int path_id = i % N_PATHS;\n",
    "    float s_curr = S0[batch_id];\n",
    "    float tmp1 = mu[batch_id]*T/N_STEPS;\n",
    "    float tmp2 = exp(-r[batch_id]*T);\n",
    "    unsigned n=0;\n",
    "    double running_average = 0.0;\n",
    "    for(unsigned n = 0; n < N_STEPS; n++){\n",
    "       s_curr += tmp1 * s_curr + sigma[batch_id]*s_curr*tmp3*d_normals[path_id + batch_id * N_PATHS + n * N_PATHS * N_BATCH];\n",
    "       running_average += (s_curr - running_average) / (n + 1.0);\n",
    "       if (running_average <= B[batch_id]){\n",
    "           break;\n",
    "       }\n",
    "    }\n",
    "\n",
    "    float payoff = (running_average>K[batch_id] ? running_average-K[batch_id] : 0.f); \n",
    "    d_s[i] = tmp2 * payoff;\n",
    "  }\n",
    "}\n",
    "\n",
    "''', 'batched_barrier_option')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note, the parameters (K, B, S0, sigma, mu, r) are passed in as an array with length of batch size. The output array is a two dimensional array flatten to 1-D. The first dimension is for Batch and the second dimension is for Path. \n",
    "\n",
    "Testing it out by entering two sets of option parameters:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_PATHS = 2048000\n",
    "N_STEPS = 365\n",
    "N_BATCH = 2\n",
    "T = 1.0\n",
    "\n",
    "K = cupy.array([110.0, 120.0], dtype=cupy.float32)\n",
    "B = cupy.array([100.0, 90.0], dtype=cupy.float32)\n",
    "S0 = cupy.array([120.0, 100.0], dtype=cupy.float32)\n",
    "sigma = cupy.array([0.35, 0.2], dtype=cupy.float32)\n",
    "mu = cupy.array([0.15, 0.1], dtype=cupy.float32)\n",
    "r =cupy.array([0.05, 0.05], dtype=cupy.float32)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Put everything into a simple function to launch this GPU kernel. The option prices for each batch is the average of the corresponding path terminal values. This can be computed easily by Cupy function `mean(axis=1)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time 0.013919591903686523 v [21.22405    0.8480416]\n"
     ]
    }
   ],
   "source": [
    "def batch_run():\n",
    "    number_of_threads = 256\n",
    "    number_of_blocks = (N_PATHS * N_BATCH - 1) // number_of_threads + 1\n",
    "    randoms_gpu = cupy.random.normal(0, 1, N_BATCH*N_PATHS * N_STEPS, dtype=cupy.float32)\n",
    "    output = cupy.zeros(N_BATCH*N_PATHS, dtype=cupy.float32)\n",
    "    cupy.cuda.stream.get_current_stream().synchronize()\n",
    "    s = time.time()\n",
    "    cupy_batched_barrier_option((number_of_blocks,), (number_of_threads,),\n",
    "                       (output, np.float32(T), K, B, S0, sigma, mu, r,\n",
    "                        randoms_gpu, N_STEPS, N_PATHS, N_BATCH))\n",
    "    v = output.reshape(N_BATCH, N_PATHS).mean(axis=1)\n",
    "    cupy.cuda.stream.get_current_stream().synchronize()\n",
    "    e = time.time()\n",
    "    print('time', e-s, 'v',v)\n",
    "batch_run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This produces the option prices $21.22$ and $0.848$ for these two sets of option parameters in $66ms$.\n",
    "\n",
    "It works efficiently hence we will construct an `OptionDataSet` class to wrap the above code so we can use it in Pytorch. For every `next` element, it generates uniform distributed random option parameters in the specified range, launches the GPU kernel to compute the option prices, convert the CuPy array to Pytorch tensors with zero copy via the DLPack. Note how we implemented the iterable Dataset interface:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OptionDataSet(torch.utils.data.IterableDataset):\n",
    "    \n",
    "    def __init__(self, max_len=10, number_path = 1000, batch=2, threads=256,seed=15):\n",
    "        self.num = 0\n",
    "        self.max_length = max_len\n",
    "        self.N_PATHS = number_path\n",
    "        self.N_STEPS = 365\n",
    "        self.N_BATCH = batch\n",
    "        self.T = np.float32(1.0)\n",
    "        self.output = cupy.zeros(self.N_BATCH*self.N_PATHS, dtype=cupy.float32) \n",
    "        self.number_of_blocks = (self.N_PATHS * self.N_BATCH - 1) // threads + 1\n",
    "        self.number_of_threads = threads\n",
    "        cupy.random.seed(seed)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.max_length\n",
    "        \n",
    "    def __iter__(self):\n",
    "        self.num = 0\n",
    "        return self\n",
    "    \n",
    "    def __next__(self):\n",
    "        if self.num > self.max_length:\n",
    "            raise StopIteration\n",
    "        X = cupy.random.rand(self.N_BATCH, 6, dtype=cupy.float32)\n",
    "        # scale the [0, 1) random numbers to the correct range for each of the option parameters\n",
    "        X = X * cupy.array([200.0, 0.99, 200.0, 0.4, 0.2, 0.2], dtype=cupy.float32)\n",
    "        # make sure the Barrier is smaller than the Strike price\n",
    "        X[:, 1] = X[:, 0] * X[:, 1]\n",
    "        randoms = cupy.random.normal(0, 1, self.N_BATCH * self.N_PATHS * self.N_STEPS, dtype=cupy.float32)\n",
    "        cupy_batched_barrier_option((self.number_of_blocks,), (self.number_of_threads,), (self.output, self.T, cupy.ascontiguousarray(X[:, 0]), \n",
    "                              cupy.ascontiguousarray(X[:, 1]), cupy.ascontiguousarray(X[:, 2]), cupy.ascontiguousarray(X[:, 3]), cupy.ascontiguousarray(X[:, 4]), cupy.ascontiguousarray(X[:, 5]), randoms, self.N_STEPS, self.N_PATHS, self.N_BATCH))\n",
    "        Y = self.output.reshape(self.N_BATCH, self.N_PATHS).mean(axis=1)\n",
    "        self.num += 1\n",
    "        return (from_dlpack(X.toDlpack()), from_dlpack(Y.toDlpack()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Put everything related to Pytorch dataset into a file `cupy_dataset.py`:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting cupy_dataset.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile cupy_dataset.py \n",
    "import cupy\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.dlpack import from_dlpack\n",
    "cupy.cuda.set_allocator(None)\n",
    "\n",
    "cupy_batched_barrier_option = cupy.RawKernel(r'''\n",
    "extern \"C\" __global__ void batched_barrier_option(\n",
    "    float *d_s,\n",
    "    const float T,\n",
    "    const float * K,\n",
    "    const float * B,\n",
    "    const float * S0,\n",
    "    const float * sigma,\n",
    "    const float * mu,\n",
    "    const float * r,\n",
    "    const float * d_normals,\n",
    "    const long N_STEPS,\n",
    "    const long N_PATHS,\n",
    "    const long N_BATCH)\n",
    "{\n",
    "  unsigned idx =  threadIdx.x + blockIdx.x * blockDim.x;\n",
    "  unsigned stride = blockDim.x * gridDim.x;\n",
    "  unsigned tid = threadIdx.x;\n",
    "  const float tmp3 = sqrt(T/N_STEPS);\n",
    "\n",
    "\n",
    "  for (unsigned i = idx; i<N_PATHS * N_BATCH; i+=stride)\n",
    "  {\n",
    "    int batch_id = i / N_PATHS;\n",
    "    int path_id = i % N_PATHS;\n",
    "    float s_curr = S0[batch_id];\n",
    "    float tmp1 = mu[batch_id]*T/N_STEPS;\n",
    "    float tmp2 = exp(-r[batch_id]*T);\n",
    "    unsigned n=0;\n",
    "    double running_average = 0.0;\n",
    "    for(unsigned n = 0; n < N_STEPS; n++){\n",
    "       s_curr += tmp1 * s_curr + sigma[batch_id]*s_curr*tmp3*d_normals[path_id + batch_id * N_PATHS + n * N_PATHS * N_BATCH];\n",
    "       running_average += (s_curr - running_average) / (n + 1.0);\n",
    "       if (running_average <= B[batch_id]){\n",
    "           break;\n",
    "       }\n",
    "    }\n",
    "\n",
    "    float payoff = (running_average>K[batch_id] ? running_average-K[batch_id] : 0.f); \n",
    "    d_s[i] = tmp2 * payoff;\n",
    "  }\n",
    "}\n",
    "\n",
    "''', 'batched_barrier_option')\n",
    "\n",
    "class OptionDataSet(torch.utils.data.IterableDataset):\n",
    "    \n",
    "    def __init__(self, max_len=10, number_path = 1000, batch=2, threads=256,seed=15):\n",
    "        self.num = 0\n",
    "        self.max_length = max_len\n",
    "        self.N_PATHS = number_path\n",
    "        self.N_STEPS = 365\n",
    "        self.N_BATCH = batch\n",
    "        self.T = np.float32(1.0)\n",
    "        self.output = cupy.zeros(self.N_BATCH*self.N_PATHS, dtype=cupy.float32) \n",
    "        self.number_of_blocks = (self.N_PATHS * self.N_BATCH - 1) // threads + 1\n",
    "        self.number_of_threads = threads\n",
    "        cupy.random.seed(seed)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.max_length\n",
    "        \n",
    "    def __iter__(self):\n",
    "        self.num = 0\n",
    "        return self\n",
    "    \n",
    "    def __next__(self):\n",
    "        if self.num > self.max_length:\n",
    "            raise StopIteration\n",
    "        X = cupy.random.rand(self.N_BATCH, 6, dtype=cupy.float32)\n",
    "        # scale the [0, 1) random numbers to the correct range for each of the option parameters\n",
    "        X = X * cupy.array([200.0, 0.99, 200.0, 0.4, 0.2, 0.2], dtype=cupy.float32)\n",
    "        # make sure the Barrier is smaller than the Strike price\n",
    "        X[:, 1] = X[:, 0] * X[:, 1]\n",
    "        randoms = cupy.random.normal(0, 1, self.N_BATCH * self.N_PATHS * self.N_STEPS, dtype=cupy.float32)\n",
    "        cupy_batched_barrier_option((self.number_of_blocks,), (self.number_of_threads,), (self.output, self.T, cupy.ascontiguousarray(X[:, 0]), \n",
    "                              cupy.ascontiguousarray(X[:, 1]), cupy.ascontiguousarray(X[:, 2]), cupy.ascontiguousarray(X[:, 3]), cupy.ascontiguousarray(X[:, 4]), cupy.ascontiguousarray(X[:, 5]), randoms, self.N_STEPS, self.N_PATHS, self.N_BATCH))\n",
    "        Y = self.output.reshape(self.N_BATCH, self.N_PATHS).mean(axis=1)\n",
    "        self.num += 1\n",
    "        return (from_dlpack(X.toDlpack()), from_dlpack(Y.toDlpack()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a test code to sample 10 data points with batch size 16:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.6558e+02, 0.0000e+00, 8.0069e+01, 1.0866e+02, 7.7740e-03, 0.0000e+00,\n",
      "        2.7772e+01, 0.0000e+00, 0.0000e+00, 6.4279e+01, 0.0000e+00, 5.1346e+00,\n",
      "        0.0000e+00, 1.4733e+02, 4.1851e+01, 0.0000e+00], device='cuda:0')\n",
      "tensor([ 57.1285,   0.0000,   0.0000, 151.9433,   0.0000,   0.0000,   0.0000,\n",
      "          9.3306,   0.0000,   0.7246, 157.0885,  10.7096,   0.0000,   0.7067,\n",
      "         59.1110,  14.6442], device='cuda:0')\n",
      "tensor([106.4531,   0.0000,  51.1248,  12.7823,  67.4821,   0.0000,   7.3539,\n",
      "          0.0000, 143.2203,  66.0655,  66.5477, 129.6811,   0.0000,  13.5559,\n",
      "         27.5546,   0.0000], device='cuda:0')\n",
      "tensor([4.1777e+01, 0.0000e+00, 2.5890e+00, 1.4500e+02, 0.0000e+00, 1.5099e+00,\n",
      "        1.1183e+02, 5.6967e+01, 7.5750e-05, 1.2390e+01, 0.0000e+00, 3.0183e+01,\n",
      "        1.3890e+01, 5.0533e+01, 3.8499e+01, 8.2232e+01], device='cuda:0')\n",
      "tensor([1.0687e+02, 3.0590e+01, 8.5428e+01, 1.9835e+01, 3.0602e+01, 1.5230e+00,\n",
      "        0.0000e+00, 0.0000e+00, 4.0244e+01, 0.0000e+00, 3.7487e-01, 0.0000e+00,\n",
      "        1.1777e+02, 0.0000e+00, 9.6200e+00, 4.2073e-04], device='cuda:0')\n",
      "tensor([ 83.6088, 125.8481,   0.0000,   0.0000,   0.0000,  35.1237,  26.4887,\n",
      "        114.6908,   1.2338, 133.6484,  84.3443,  49.0381,  33.3620,  93.0905,\n",
      "         40.8572,  30.2684], device='cuda:0')\n",
      "tensor([1.6068e+01, 6.8251e+01, 1.7516e+00, 6.3889e+01, 2.0682e+00, 3.0282e-01,\n",
      "        2.3074e-04, 2.4942e+01, 1.1639e+02, 0.0000e+00, 3.0597e+01, 0.0000e+00,\n",
      "        3.0390e+01, 2.1144e+00, 8.2769e-04, 6.3105e+01], device='cuda:0')\n",
      "tensor([129.0360,   0.0000,   0.0000,  34.7129,  76.3240,  61.5014,  96.1047,\n",
      "         41.5991,   0.0000,   0.0000,   1.6868,   0.0000,   0.0000, 198.8765,\n",
      "          0.0000, 130.8935], device='cuda:0')\n",
      "tensor([23.4824, 49.1953, 70.5731,  0.0000,  0.0000, 35.5231,  0.0000,  0.0000,\n",
      "         0.0000, 64.7130,  0.0000, 56.6821,  3.6377,  0.0000,  0.0000, 17.6415],\n",
      "       device='cuda:0')\n",
      "tensor([113.4123,   0.2840,   0.0000,   9.8790,  34.9789,  62.0461,   0.0000,\n",
      "          0.0000,  90.4281, 151.8807,   0.0000,   0.0000,  75.6426, 137.9153,\n",
      "          0.0000,  65.4237], device='cuda:0')\n",
      "tensor([1.1853e+02, 0.0000e+00, 0.0000e+00, 3.5182e+01, 8.2466e+01, 0.0000e+00,\n",
      "        0.0000e+00, 1.7089e+01, 0.0000e+00, 2.8777e-02, 0.0000e+00, 0.0000e+00,\n",
      "        6.7766e+01, 3.9360e+01, 1.2019e+02, 1.0623e+02], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "from cupy_dataset import OptionDataSet\n",
    "ds = OptionDataSet(10, number_path=100000, batch=16, seed=15)\n",
    "for i in ds:\n",
    "    print(i[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can implement the same code by using Numba to accelerate the calculation in GPU:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.6558e+02, 0.0000e+00, 8.0069e+01, 1.0866e+02, 7.7740e-03, 0.0000e+00,\n",
      "        2.7772e+01, 0.0000e+00, 0.0000e+00, 6.4279e+01, 0.0000e+00, 5.1346e+00,\n",
      "        0.0000e+00, 1.4733e+02, 4.1851e+01, 0.0000e+00], device='cuda:0')\n",
      "tensor([ 57.1285,   0.0000,   0.0000, 151.9433,   0.0000,   0.0000,   0.0000,\n",
      "          9.3306,   0.0000,   0.7246, 157.0885,  10.7096,   0.0000,   0.7067,\n",
      "         59.1110,  14.6442], device='cuda:0')\n",
      "tensor([106.4531,   0.0000,  51.1248,  12.7823,  67.4821,   0.0000,   7.3539,\n",
      "          0.0000, 143.2203,  66.0655,  66.5476, 129.6811,   0.0000,  13.5559,\n",
      "         27.5546,   0.0000], device='cuda:0')\n",
      "tensor([4.1777e+01, 0.0000e+00, 2.5890e+00, 1.4500e+02, 0.0000e+00, 1.5099e+00,\n",
      "        1.1183e+02, 5.6967e+01, 7.5751e-05, 1.2390e+01, 0.0000e+00, 3.0183e+01,\n",
      "        1.3890e+01, 5.0533e+01, 3.8499e+01, 8.2232e+01], device='cuda:0')\n",
      "tensor([1.0687e+02, 3.0590e+01, 8.5428e+01, 1.9835e+01, 3.0602e+01, 1.5230e+00,\n",
      "        0.0000e+00, 0.0000e+00, 4.0244e+01, 0.0000e+00, 3.7487e-01, 0.0000e+00,\n",
      "        1.1777e+02, 0.0000e+00, 9.6200e+00, 4.2073e-04], device='cuda:0')\n",
      "tensor([ 83.6088, 125.8481,   0.0000,   0.0000,   0.0000,  35.1237,  26.4887,\n",
      "        114.6908,   1.2338, 133.6484,  84.3443,  49.0380,  33.3620,  93.0905,\n",
      "         40.8572,  30.2684], device='cuda:0')\n",
      "tensor([1.6068e+01, 6.8251e+01, 1.7516e+00, 6.3889e+01, 2.0682e+00, 3.0282e-01,\n",
      "        2.3074e-04, 2.4942e+01, 1.1639e+02, 0.0000e+00, 3.0597e+01, 0.0000e+00,\n",
      "        3.0390e+01, 2.1144e+00, 8.2769e-04, 6.3105e+01], device='cuda:0')\n",
      "tensor([129.0360,   0.0000,   0.0000,  34.7129,  76.3240,  61.5014,  96.1047,\n",
      "         41.5991,   0.0000,   0.0000,   1.6868,   0.0000,   0.0000, 198.8765,\n",
      "          0.0000, 130.8935], device='cuda:0')\n",
      "tensor([23.4824, 49.1953, 70.5731,  0.0000,  0.0000, 35.5231,  0.0000,  0.0000,\n",
      "         0.0000, 64.7129,  0.0000, 56.6821,  3.6377,  0.0000,  0.0000, 17.6415],\n",
      "       device='cuda:0')\n",
      "tensor([113.4123,   0.2840,   0.0000,   9.8790,  34.9789,  62.0461,   0.0000,\n",
      "          0.0000,  90.4281, 151.8807,   0.0000,   0.0000,  75.6426, 137.9153,\n",
      "          0.0000,  65.4237], device='cuda:0')\n",
      "tensor([1.1853e+02, 0.0000e+00, 0.0000e+00, 3.5182e+01, 8.2466e+01, 0.0000e+00,\n",
      "        0.0000e+00, 1.7089e+01, 0.0000e+00, 2.8777e-02, 0.0000e+00, 0.0000e+00,\n",
      "        6.7766e+01, 3.9360e+01, 1.2019e+02, 1.0623e+02], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "import numba\n",
    "from numba import cuda\n",
    "\n",
    "@cuda.jit\n",
    "def batch_barrier_option(d_s, T, K, B, S0, sigma, mu, r, d_normals, N_STEPS, N_PATHS, N_BATCH):\n",
    "    # ii - overall thread index\n",
    "    ii = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x\n",
    "    stride = cuda.gridDim.x * cuda.blockDim.x\n",
    "    tmp3 = math.sqrt(T/N_STEPS)\n",
    "    for i in range(ii, N_PATHS * N_BATCH, stride):\n",
    "        batch_id = i // N_PATHS\n",
    "        path_id = i % N_PATHS\n",
    "        tmp1 = mu[batch_id]*T/N_STEPS\n",
    "        tmp2 = math.exp(-r[batch_id]*T)\n",
    "        running_average = 0.0\n",
    "        s_curr = S0[batch_id]\n",
    "        for n in range(N_STEPS):\n",
    "\n",
    "            s_curr += tmp1 * s_curr + sigma[batch_id]*s_curr*tmp3*d_normals[path_id + batch_id * N_PATHS + n * N_PATHS * N_BATCH]\n",
    "            running_average = running_average + 1.0/(n + 1.0) * (s_curr - running_average)\n",
    "            if i==0 and batch_id == 2:\n",
    "                print(s_curr)\n",
    "            if running_average <= B[batch_id]:\n",
    "                break\n",
    "        payoff = running_average - K[batch_id] if running_average > K[batch_id] else 0\n",
    "        d_s[i] = tmp2 * payoff\n",
    "\n",
    "class NumbaOptionDataSet(object):\n",
    "    \n",
    "    def __init__(self, max_len=10, number_path = 1000, batch=2, threads=512, seed=15):\n",
    "        self.num = 0\n",
    "        self.max_length = max_len\n",
    "        self.N_PATHS = number_path\n",
    "        self.N_STEPS = 365\n",
    "        self.N_BATCH = batch\n",
    "        self.T = np.float32(1.0)\n",
    "        self.output = cupy.zeros(self.N_BATCH*self.N_PATHS, dtype=cupy.float32) \n",
    "        self.number_of_blocks = (self.N_PATHS * self.N_BATCH - 1) // threads + 1\n",
    "        self.number_of_threads = threads\n",
    "        cupy.random.seed(seed)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.max_length\n",
    "        \n",
    "    def __iter__(self):\n",
    "        self.num = 0\n",
    "        return self\n",
    "    \n",
    "    def __next__(self):\n",
    "        if self.num > self.max_length:\n",
    "            raise StopIteration\n",
    "        X = cupy.random.rand(self.N_BATCH, 6, dtype=cupy.float32)\n",
    "        # scale the [0, 1) random numbers to the correct range for each of the option parameters\n",
    "        X = X * cupy.array([200.0, 0.99, 200.0, 0.4, 0.2, 0.2], dtype=cupy.float32)\n",
    "        # make sure the Barrier is smaller than the Strike price\n",
    "        X[:, 1] = X[:, 0] * X[:, 1]\n",
    "        randoms = cupy.random.normal(0, 1, self.N_BATCH * self.N_PATHS * self.N_STEPS, dtype=cupy.float32)\n",
    "        batch_barrier_option[(self.number_of_blocks,), (self.number_of_threads,)](self.output, self.T, X[:, 0], \n",
    "                              X[:, 1], X[:, 2], X[:, 3], X[:, 4], X[:, 5], randoms, self.N_STEPS, self.N_PATHS, self.N_BATCH)\n",
    "        o = self.output.reshape(self.N_BATCH, self.N_PATHS)\n",
    "        Y = o.mean(axis = 1) \n",
    "        self.num += 1\n",
    "        return (from_dlpack(X.toDlpack()), from_dlpack(Y.toDlpack()))\n",
    "ds = NumbaOptionDataSet(10, number_path=100000, batch=16, seed=15)\n",
    "for i in ds:\n",
    "    print(i[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model\n",
    "To map the option parameters to price, we use 6 layers of fully connected neural network with hidden dimension 512 as inspired by [this paper](https://arxiv.org/abs/1809.02233). Writing this DL price model into a file `model.py`:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting model.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile model.py\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "\n",
    "    def __init__(self, hidden=1024):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(6, hidden)\n",
    "        self.fc2 = nn.Linear(hidden, hidden)\n",
    "        self.fc3 = nn.Linear(hidden, hidden)\n",
    "        self.fc4 = nn.Linear(hidden, hidden)\n",
    "        self.fc5 = nn.Linear(hidden, hidden)\n",
    "        self.fc6 = nn.Linear(hidden, 1)\n",
    "        self.register_buffer('norm',\n",
    "                             torch.tensor([200.0,\n",
    "                                           198.0,\n",
    "                                           200.0,\n",
    "                                           0.4,\n",
    "                                           0.2,\n",
    "                                           0.2]))\n",
    "\n",
    "    def forward(self, x):\n",
    "        # normalize the parameter to range [0-1] \n",
    "        x = x / self.norm\n",
    "        x = F.elu(self.fc1(x))\n",
    "        x = F.elu(self.fc2(x))\n",
    "        x = F.elu(self.fc3(x))\n",
    "        x = F.elu(self.fc4(x))\n",
    "        x = F.elu(self.fc5(x))\n",
    "        return self.fc6(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we know the random parameters' scaling factors, the input parameters are first scaled back to a range of (0-1) by dividing them by (200.0, 198.0, 200.0, 0.4, 0.2, 0.2). Then they are projected 5 times to the hidden dimension of 512 after the `ELu` activation function. `ELu` is chosen because we need to compute the second order differentiation of the parameters. If use ReLu, the second order differentiation will always be zero. The last layer is a linear layer that maps the hidden dimension to the predicted option price. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For training, we use [Ignite](https://github.com/pytorch/ignite) which is a high-level library to train neural networks in PyTorch. We use `MSELoss` as the loss function, `Adam` as the optimizer and `CosineAnnealingScheduler` as the learning rate scheduler. The following code is feeding the random option data to the pricing model to train it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ignite.engine import Engine, Events\n",
    "from ignite.handlers import Timer\n",
    "from torch.nn import MSELoss\n",
    "from torch.optim import Adam\n",
    "from ignite.contrib.handlers.param_scheduler import CosineAnnealingScheduler\n",
    "from ignite.handlers import ModelCheckpoint\n",
    "from model import Net\n",
    "from cupy_dataset import OptionDataSet\n",
    "timer = Timer(average=True)\n",
    "model = Net().cuda()\n",
    "loss_fn = MSELoss()\n",
    "optimizer = Adam(model.parameters(), lr=1e-3)\n",
    "dataset = OptionDataSet(max_len=10000, number_path = 1024, batch=4800)\n",
    "\n",
    "def train_update(engine, batch):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    x = batch[0]\n",
    "    y = batch[1]\n",
    "    y_pred = model(x)\n",
    "    loss = loss_fn(y_pred[:,0], y)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    return loss.item()\n",
    "\n",
    "trainer = Engine(train_update)\n",
    "log_interval = 100\n",
    "\n",
    "scheduler = CosineAnnealingScheduler(optimizer, 'lr', 1e-4, 1e-6, len(dataset))\n",
    "trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)\n",
    "timer.attach(trainer,\n",
    "             start=Events.EPOCH_STARTED,\n",
    "             resume=Events.ITERATION_STARTED,\n",
    "             pause=Events.ITERATION_COMPLETED,\n",
    "             step=Events.ITERATION_COMPLETED)    \n",
    "@trainer.on(Events.ITERATION_COMPLETED)\n",
    "def log_training_loss(engine):\n",
    "    iter = (engine.state.iteration - 1) % len(dataset) + 1\n",
    "    if iter % log_interval == 0:\n",
    "        print('loss', engine.state.output, 'average time', timer.value())\n",
    "        \n",
    "trainer.run(dataset, max_epochs=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss is keeping decreasing which means the pricing model can predict the option prices better. It takes about $12ms$ to compute one mini-batch in average, In the following sections, we will try to expore the full potentials of the GPU to accelerate the training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TensorCore mixed precision training\n",
    "\n",
    "The V100 GPUs have 640 tensor cores that can accelerate half precision matrix multiplication calculation which is the core computation done by the DL model. [Apex library](https://github.com/NVIDIA/apex) developed by NVIDIA makes mixed precision and distributed training in Pytorch easy. By changing 3 lines of code, it can use the tensor cores to accelerate the training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from apex import amp\n",
    "from ignite.engine import Engine, Events\n",
    "from torch.nn import MSELoss\n",
    "from ignite.handlers import Timer\n",
    "from torch.optim import Adam\n",
    "from ignite.contrib.handlers.param_scheduler import CosineAnnealingScheduler\n",
    "from ignite.handlers import ModelCheckpoint\n",
    "from model import Net\n",
    "from cupy_dataset import OptionDataSet\n",
    "timer = Timer(average=True)\n",
    "model = Net().cuda()\n",
    "loss_fn = MSELoss()\n",
    "optimizer = Adam(model.parameters(), lr=1e-3)\n",
    "# set the AMP optimization level to O1\n",
    "opt_level = 'O1'\n",
    "# wrap the optimizer and model\n",
    "model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)\n",
    "dataset = OptionDataSet(max_len=10000, number_path = 1024, batch=4800)\n",
    "\n",
    "def train_update(engine, batch):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    x = batch[0]\n",
    "    y = batch[1]\n",
    "    y_pred = model(x)\n",
    "    loss = loss_fn(y_pred[:,0], y)\n",
    "    # amp handles the auto loss scaling\n",
    "    with amp.scale_loss(loss, optimizer) as scaled_loss:\n",
    "        scaled_loss.backward()\n",
    "    optimizer.step()\n",
    "    return loss.item()\n",
    "\n",
    "trainer = Engine(train_update)\n",
    "log_interval = 100\n",
    "timer.attach(trainer,\n",
    "             start=Events.EPOCH_STARTED,\n",
    "             resume=Events.ITERATION_STARTED,\n",
    "             pause=Events.ITERATION_COMPLETED,\n",
    "             step=Events.ITERATION_COMPLETED)    \n",
    "scheduler = CosineAnnealingScheduler(optimizer, 'lr', 1e-4, 1e-6, len(dataset))\n",
    "trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)\n",
    "    \n",
    "@trainer.on(Events.ITERATION_COMPLETED)\n",
    "def log_training_loss(engine):\n",
    "    iter = (engine.state.iteration - 1) % len(dataset) + 1\n",
    "    if iter % log_interval == 0:\n",
    "        print('loss', engine.state.output, 'average time', timer.value())\n",
    "        \n",
    "trainer.run(dataset, max_epochs=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It improves to compute each mini-batch in $8ms$. As we reduce the model weights to half precision for better performance, the loss need to be scaled to make sure the half precision dynamic range aligns with the computation. It is guessing what is the correct loss scaling factor and adjust it automatically if the gradient overflows. In the end, we will get the best hardware acceleration while maintaining the accuracy of model prediction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multiple GPU training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Apex makes multiple GPU training easy. Working on the same training script, we need to take care of a few extra steps:\n",
    "\n",
    "1. Add the argument `--local_rank` which will be automatically set by the distributed launcher\n",
    "2. Initialize the process group\n",
    "2. Generate independent batched data based on process id in the dataset.\n",
    "3. Wrap the model and optimizer to handle distributed computation. \n",
    "4. Scale the loss and optimizer\n",
    "\n",
    "To launch distributed training, we need to put everything into a python file. Following is an example:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting distributed_train.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile distributed_train.py \n",
    "import cupy\n",
    "import numpy as np\n",
    "import math\n",
    "import time\n",
    "import os\n",
    "import torch\n",
    "from torch.utils.dlpack import from_dlpack\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "from apex import amp\n",
    "from ignite.engine import Engine, Events\n",
    "from torch.nn import MSELoss\n",
    "from torch.optim import Adam\n",
    "from ignite.contrib.handlers.param_scheduler import CosineAnnealingScheduler\n",
    "from ignite.handlers import ModelCheckpoint\n",
    "from apex.parallel import DistributedDataParallel \n",
    "import argparse\n",
    "from model import Net\n",
    "from cupy_dataset import OptionDataSet\n",
    "\n",
    "parser = argparse.ArgumentParser()\n",
    "parser = argparse.ArgumentParser()\n",
    "# this local_rank arg is automaticall set by distributed launch\n",
    "parser.add_argument(\"--local_rank\", default=0, type=int)\n",
    "args = parser.parse_args()\n",
    "\n",
    "args.distributed = False\n",
    "if 'WORLD_SIZE' in os.environ:\n",
    "    args.distributed = int(os.environ['WORLD_SIZE']) > 1\n",
    "\n",
    "if args.distributed:\n",
    "    torch.cuda.set_device(args.local_rank)\n",
    "    torch.distributed.init_process_group(backend='nccl',\n",
    "                                         init_method='env://')\n",
    "\n",
    "torch.backends.cudnn.benchmark = True\n",
    "\n",
    "\n",
    "model = Net().cuda()\n",
    "loss_fn = MSELoss()\n",
    "optimizer = Adam(model.parameters(), lr=1e-3)\n",
    "opt_level = 'O1'\n",
    "model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)\n",
    "if args.distributed:\n",
    "    model = DistributedDataParallel(model)\n",
    "dataset = OptionDataSet(max_len=10000, number_path = 1024, batch=10240, seed=args.local_rank)\n",
    "\n",
    "def train_update(engine, batch):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    x = batch[0]\n",
    "    y = batch[1]\n",
    "    y_pred = model(x)\n",
    "    loss = loss_fn(y_pred[:,0], y)\n",
    "    with amp.scale_loss(loss, optimizer) as scaled_loss:\n",
    "        scaled_loss.backward()\n",
    "    optimizer.step()\n",
    "    return loss.item()\n",
    "\n",
    "trainer = Engine(train_update)\n",
    "log_interval = 100\n",
    "\n",
    "scheduler = CosineAnnealingScheduler(optimizer, 'lr', 1e-4, 1e-6, len(dataset))\n",
    "trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)\n",
    "    \n",
    "@trainer.on(Events.ITERATION_COMPLETED)\n",
    "def log_training_loss(engine):\n",
    "    iter = (engine.state.iteration - 1) % len(dataset) + 1\n",
    "    if iter % log_interval == 0:\n",
    "        print('loss', engine.state.output)\n",
    "        \n",
    "trainer.run(dataset, max_epochs=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To launch multiple processes training, we need to run the following command:-"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reset -f\n",
    "\n",
    "!python -m torch.distributed.launch --nproc_per_node=4 distributed_train.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It works and all the GPUs are busy to train this network. However, it has a few problems:-\n",
    "   \n",
    "    1. There is no model serialization so the trained model is not saved\n",
    "    2. There is no validation dataset to check the training progress\n",
    "    3. Most of the time is spent in Monte Carlo simulation hence the training is slow\n",
    "    4. We use a few paths(1024) for each option parameter set which is noise and the model cannot converge to a low cost value.\n",
    "We will address these problems in the next notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
